package code.jzh.monitor.service;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import code.jzh.monitor.util.AgentConstants;
import code.jzh.monitor.util.PropertiesUtil;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import javassist.NotFoundException;
/**
 * 代理类对允许的方法字节码增强
 * @author shuyingwu
 *
 */
public class MonitorTransformer implements ClassFileTransformer {

	final static String prefix = "\nlong startTime = System.currentTimeMillis();";
	final static String postfix = "\nlong endTime = System.currentTimeMillis();";
	final static List<String> clazzRegList = PropertiesUtil.getRegList(AgentConstants.REG_KEY);
	final static List<String> methodRegList = PropertiesUtil.getRegList(AgentConstants.REG_METHOD);
	final static List<String> methodUnRegList = PropertiesUtil.getRegList(AgentConstants.UN_REG_METHOD);

	private final static Set<String> ClassSet = new HashSet<String>();
	private final static Set<String> mothodSet = new HashSet<String>();

	@Override
	public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
			ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
		if (MonitorTransformer.ClassSet.contains(className)) {
			return null;
		}
		MonitorTransformer.ClassSet.add(className);
		// 只拦截配置过需要添加字节码的类，通过instrumentation进来的class路径用‘/’分割
		for (String ckass : clazzRegList) {
			// 将路径中‘/’替换为‘.’
			className = className.replace("/", ".");
			if (className.equals(ckass)) {
				CtClass ctclass = null;
				ClassPool pool = ClassPool.getDefault();
				try {
					// 系统类加载器获取不到tomcat，jboss等自定义类加载器加载的类，所以需要把他们的类路径配置进去
					pool.appendClassPath(new LoaderClassPath(loader));
					// 用于取得字节码类，必须在当前的classpath中，使用全称 ,这部分是关于javassist的知识
					ctclass = pool.get(className);
					CtMethod[] ctmethodArr = ctclass.getDeclaredMethods();
					// 循环一下，看看哪些方法需要加时间监测
					for (CtMethod ctmethod : ctmethodArr) {
						String method = ctmethod.getLongName();
						// 获取方法名
						if (!ctmethod.isEmpty() && !methodUnRegList.contains(method.substring(0, method.indexOf("(")))
								&& (methodRegList.contains(ckass + ".*")
										|| methodRegList.contains(method.substring(0, method.indexOf("("))))) {
							if (MonitorTransformer.mothodSet.contains(method)) {
								break;
							}
							MonitorTransformer.mothodSet.add(method);
							StringBuffer insertBeforeBuf = new StringBuffer();
							insertBeforeBuf.append("{").append(prefix)
									.append("code.jzh.monitor.util.TraceJsonUtil.getTraceJosnStr(startTime,\"" + className
											+ "\",\"" + ctmethod.getName() + "\","+AgentConstants.METHOD_START+",0);\n")
									.append("}");
							ctmethod.insertBefore(insertBeforeBuf.toString());
							StringBuffer insertAfterBuf = new StringBuffer();
							insertAfterBuf.append("{").append(postfix)
									.append("code.jzh.monitor.util.TraceJsonUtil.getTraceJosnStr(endTime,\"" + className
											+ "\",\"" + ctmethod.getName() + "\","+AgentConstants.METHOD_END+",0);\n")
									.append("}");
							ctmethod.insertAfter(insertAfterBuf.toString());
							CtClass etype = ClassPool.getDefault().get("java.lang.Exception");
							StringBuffer catchBuff = new StringBuffer();
							catchBuff.append("{").append(postfix)
									 .append("code.jzh.monitor.util.TraceJsonUtil.getTraceJosnStr(endTime,\"" + className
											 + "\",\"" + ctmethod.getName() + "\","+AgentConstants.METHOD_END+",1);\n")
									 .append("throw $e;}");
							ctmethod.addCatch(catchBuff.toString(), etype);
						}
					}
					return ctclass.toBytecode();
				} catch (IOException e) {
					e.printStackTrace();
				} catch (CannotCompileException e) {
					e.printStackTrace();
				} catch (NotFoundException e) {
					e.printStackTrace();
				}
			}
		}
		return null;
	}

}